import lightgbm as lgb
import pandas as pd
import backtrader as bt
import pyfolio as pf


# 加载模型
model_path = "model/cls-20_prft-0.03_quanA_label-profit_30fac_2024-12-02.txt"
model = lgb.Booster(model_file=model_path)

# 加载因子数据
factor_data = pd.read_csv("2023_factors.csv", parse_dates=["datetime"])
factor_data.set_index("datetime", inplace=True)

# 加载行情数据（例如收盘价）
price_data = pd.read_csv("2023_prices.csv", parse_dates=["datetime"])
price_data.set_index("datetime", inplace=True)


class LightGBMStrategy(bt.Strategy):
    params = (
        ("model", None),  # LightGBM 模型
        ("features", []), # 因子列名
    )

    def __init__(self):
        self.model = self.params.model
        self.features = self.params.features

    def next(self):
        # 获取当前时间点的因子数据
        current_date = self.data.datetime.date(0)
        try:
            factor_row = factor_data.loc[current_date].to_frame().T
            feature_values = factor_row[self.features].values
        except KeyError:
            return  # 如果没有因子数据，则跳过

        # 使用 LightGBM 模型进行预测
        prediction = self.model.predict(feature_values)[0]

        # 根据预测结果生成交易信号
        if prediction > 0:  # 假设预测值 > 0 表示买入信号
            self.buy()
        elif prediction < 0:  # 假设预测值 < 0 表示卖出信号
            self.sell()


# 创建 Cerebro 引擎
cerebro = bt.Cerebro()

# 添加数据源
data_feed = bt.feeds.PandasData(dataname=price_data)
cerebro.adddata(data_feed)

# 添加策略
cerebro.addstrategy(LightGBMStrategy, model=model, features=factor_data.columns.tolist())

# 设置初始资金
cerebro.broker.set_cash(100000.0)

# 添加分析器
cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe")
cerebro.addanalyzer(bt.analyzers.Returns, _name="returns")
cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")

# 运行回测
results = cerebro.run()

# 输出结果
final_portfolio_value = cerebro.broker.getvalue()
print(f"Final Portfolio Value: {final_portfolio_value}")
print(f"Sharpe Ratio: {results[0].analyzers.sharpe.get_analysis()['sharperatio']}")
print(f"Max Drawdown: {results[0].analyzers.drawdown.get_analysis()['max']['drawdown']}%")

# 提取回测结果
returns, positions, transactions = cerebro.runstrats[0][0].analyzers.pyfolio.get_pf_items()

# 使用 pyfolio 绘制性能报告
pf.create_full_tear_sheet(returns)